/* Copyright (c) 2025 LiuZiyao
*    Key Laboratory of Cyberspace Security Defense,Institute of Information Engineering, CAS
*    School of Cyber Security, University of Chinese Academy of Sciences     
*
* This file is part of the openHiTLS project.
*
* openHiTLS is licensed under the Mulan PSL v2.
* You can use this software according to the terms and conditions of the Mulan PSL v2.
* You may obtain a copy of Mulan PSL v2 at:
*
*     http://license.coscl.org.cn/MulanPSL2
*
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
* EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
* MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
* See the Mulan PSL v2 for more details.
*/


# include <stdio.h>
# include <stdlib.h>
# include <string.h>
#include <immintrin.h>
#include <stdint.h>
# include "polarlac_param.h"

#if defined(LAC_LIGHT)
//polar(256,128,1/2)
#include "polarlight.h"
#endif

#if defined(LAC128)
//polar(512,128,1/4)
#include "polar128.h"
#endif

#if defined(LAC256)
//polar(512,256,1/2)
#include "polar256.h"
#endif

#define sign_macro(x) ((x > 0) - (x < 0))
#define absl_macro(x) (((x > 0) - (x < 0)) * x)
#define mini_macro(x, y) ((x < y) ? x : y)
#define f_macro(L1, L2) sign_macro(L1) * sign_macro(L2) * mini_macro(absl_macro(L1), absl_macro(L2))
#define g_macro(u, L1, L2) (((1 - 2*u) * L1) + L2)

/**
 * polar encode
 * Algorithm idea inspired by:
 * https://github.com/sravan-ankireddy/polar_codes (Repository does not specify a license)
 */
void encode_polar(uint8_t *u)
{ 
    int stage_size = 1; // current butterfly size

    for (int level = 0; level < polar.n; level++)
    {
        int group_size = stage_size * 2;
        int num_groups = (1 << polar.n) / group_size;

        for (int block = 0; block < num_groups; block++)
        {
            uint8_t *segment = u + block * group_size;

            /* Apply butterfly transform */
            for (int offset = 0; offset < stage_size; offset++)
            {
                uint8_t left  = segment[offset];
                uint8_t right = segment[offset + stage_size];

                segment[offset] = left ^ right;
            }
        }

        stage_size <<= 1; // move to next level
    }
}

/**
 * polar decode
 * Algorithm idea inspired by:
 * https://github.com/YuYongRun/PolarCodeDecodersInMatlab (Repository does not specify a license)
 */
void decode_polar(uint8_t *m_cap, const double *llr)
{
    uint8_t inter_bit[2 * polar.N - 1][2];  // internal bit vector
    double inter_llr[polar.N - 1];  // internal llr vector, inter_llr[0] used for decision
    int msg_index = 0;

    for(int i = 0; i < polar.N; i++) // decode each u_i
    {
        if(i == 0)
        {
            int index1 = lambda_offset[polar.n - 1];
            int beta = 0;
            int end_beta = index1 - 1;
            
            for(; beta <= end_beta; beta++)
            {
                inter_llr[beta + index1 - 1] = f_macro(llr[beta], llr[beta + index1]);
            }

            for(int layer = polar.n - 2; layer >= -1; layer--)
            {
                int index1 = lambda_offset[layer];
                int index2 = lambda_offset[layer + 1];
                int beta = index1 - 1;
                int end_beta = index2 - 2;
                
                for(; beta <= end_beta; beta++)
                {
                    inter_llr[beta] = f_macro(inter_llr[beta + index1], inter_llr[beta + index2]);
                }
            }
        }
        else if(i == polar.N/2)
        {
            int index1 = lambda_offset[polar.n - 1];
            int beta = 0;
            int end_beta = index1 - 1;
            
            for(; beta <= end_beta; beta++)
            {
                inter_llr[beta + index1 - 1] = g_macro(inter_bit[beta + index1 - 1][0], llr[beta], llr[beta + index1]);
            }

            for(int layer = polar.n - 2; layer >= -1; layer--)
            {
                int index1 = lambda_offset[layer];
                int index2 = lambda_offset[layer + 1];
                int beta = index1 - 1;
                int end_beta = index2 - 2;
                
                for(; beta <= end_beta; beta++)
                {
                    inter_llr[beta] = f_macro(inter_llr[beta + index1], inter_llr[beta + index2]);
                }
            }
        }
        else
        {
            int llr_layer = llr_layer_vec[i];
            int index1 = lambda_offset[llr_layer];
            int index2 = lambda_offset[llr_layer + 1];
            int beta = index1 - 1;
            int end_beta = index2 - 2;
            
            for(; beta <= end_beta; beta++)
            {
                inter_llr[beta] = g_macro(inter_bit[beta][0], inter_llr[beta + index1], inter_llr[beta + index2]);
            }

            for(int layer = llr_layer - 1; layer >= -1; layer--)
            {
                int index1 = lambda_offset[layer];
                int index2 = lambda_offset[layer + 1];
                int beta = index1 - 1;
                int end_beta = index2 - 2;
                
                for(; beta <= end_beta; beta++)
                {
                    inter_llr[beta] = f_macro(inter_llr[beta + index1], inter_llr[beta + index2]);
                }  
            }
        }

        int i_mod_2 = i & 1;
        if(info_nodes[i] == 0)
        {
            inter_bit[0][i_mod_2] = 0;
        }
        else
        {
            int u_i = (inter_llr[0] < 0); // decision
            inter_bit[0][i_mod_2] = u_i;
            m_cap[msg_index] = u_i;
            msg_index++;
        }
        if(i_mod_2 == 1) // bit recursion
        {
            int bit_layer = bit_layer_vec[i];
            int index1;
            int index2;
            for(int layer = 0; layer <= bit_layer - 1; layer++)
            {
                index1 = lambda_offset[layer];
                index2 = lambda_offset[layer + 1];
                for(int beta = index1 - 1; beta <= index2 - 2; beta++)
                {
                    inter_bit[beta + index1][1] = inter_bit[beta][0] ^ inter_bit[beta][1];
                    inter_bit[beta + index2][1] = inter_bit[beta][1];
                }
            }

            index1 = lambda_offset[bit_layer];
            index2 = lambda_offset[bit_layer + 1];
            for (int beta = index1 - 1; beta <= index2 - 2; beta++)
            {
                inter_bit[beta + index1][0] = inter_bit[beta][0] ^ inter_bit[beta][1];
                inter_bit[beta + index2][0] = inter_bit[beta][1];
            }
        }
    }
}